import os  # 导入操作系统相关模块
import os.path  # 导入路径操作相关模块
import SimpleITK as sitk
import numpy as np


def nii_resize(image_path, save_folder="", standard_size=0):
    """
    将256*208或320*260的nii图像填充边界(补0),使其大小变为320*320
    图像要求：长宽尺寸均要小于等于320，且为偶数

    Args:
        image_path (string): 待修改大小的nii文件存放路径
        save_folder (string): 修改大小后的nii文件存放路径
        standard_size (int): 修改后标准图像的边长
    Returns:
        图像尺寸符合要求且文件修改成功则返回若干nii文件
    """
    if save_folder != "" and not os.path.exists(save_folder):
        os.mkdir(save_folder)  # 如果保存文件夹不存在，则创建

    for file in os.listdir(image_path):  # 遍历图像列表
        if file.endswith(".nii"):
            print('-' * 50, "\n正在处理: ", file, sep='')
            global img_itk
            try:
                img_itk = sitk.ReadImage(os.path.join(image_path, file))  # 读取医学图像
            except Exception as e:
                print('ReadImage Error:', e)
                continue

            original_size = img_itk.GetSize()  # 获取原始图像尺寸和间距
            image_array = sitk.GetArrayFromImage(img_itk)  # 获取图像数据

            need_standard_size = standard_size if standard_size != 0 else original_size[0]  # 获取设定的标准值
            pad_height = (need_standard_size - original_size[1]) // 2  # 计算填充量
            pad_weight = (need_standard_size - original_size[0]) // 2
            print("before_size: ", original_size, standard_size, pad_height, pad_weight)
            global padded_image_array
            try:
                padded_image_array = np.pad(image_array, (
                    (0, 0), (pad_height, pad_height), (pad_weight, pad_weight)
                ), mode='constant', constant_values=0)  # 使用np.pad对图像进行上下填充
            except Exception as e:
                print('padded Error:', e)
                continue

            resized_image = sitk.GetImageFromArray(padded_image_array)  # 创建新的 SimpleITK 图像

            resized_image.SetOrigin(img_itk.GetOrigin())  # 将元数据设置为与原始图像相同
            resized_image.SetSpacing(img_itk.GetSpacing())
            resized_image.SetDirection(img_itk.GetDirection())

            file_name = os.path.basename(file)
            save_path = os.path.join(save_folder, file_name)  # 生成新的文件名并保存图像
            print("before_size: ", original_size, "\nafter_size : ", resized_image.GetSize(), sep='')
            print("file_name:", save_path)

            sitk.WriteImage(resized_image, save_path)  # 保存调整大小后的 NIfTI 图像


if __name__ == '__main__':
    # 使用示例
    input_path = r"C:\Users\18049\Desktop\Good Labels\pre"
    output_path = r"/test/nii_resize_output"

    # 调用Resampling函数对医学图像进行重采样，并保存结果
    nii_resize(image_path=input_path, save_folder=output_path)

    # 2022中间过程使用
    # all_path = r"C:\Users\18049\Desktop\nii"
    # nii_path_2022 = all_path + r"\2_classification_nii\metastasis 2022"
    # nii_path_2022_resized = all_path + r"\2_classification_nii\metastasis 2022 resized"
    # for patient_nii_folder in os.listdir(nii_path_2022):
    #     patient_dcm_folder_path = os.path.join(nii_path_2022, patient_nii_folder)  # C:\..\CHEN YI JIANG_M00xxxxx
    #     if os.path.isdir(patient_dcm_folder_path):  # 检查是否是文件夹
    #         print('-' * 50, '\n', patient_nii_folder, sep='')
    #         resized_nii_folder_path = os.path.join(nii_path_2022_resized, patient_nii_folder)  # 创建存放目标nii的文件夹
    #         if not os.path.exists(resized_nii_folder_path):
    #             os.makedirs(resized_nii_folder_path)
    #         patient_nii_folder_path = os.path.join(nii_path_2022, patient_nii_folder)  # 创建存放目标nii的文件夹
    #         nii_resize(image_path=patient_nii_folder_path, save_folder=resized_nii_folder_path, standard_size=320)

# 使用重采样改变图像大小
# def nii_resize_resample(image_path, save_folder="", output_path="", label=False):
#     # 如果保存文件夹不存在，则创建
#     if save_folder != "" and not os.path.exists(save_folder):
#         os.mkdir(save_folder)
#     # 获取匹配图像路径列表
#     import glob
#     filelist = sorted(glob.glob(image_path))
#     print(filelist)
#     # 遍历图像列表
#     for file in filelist:
#         print(file)
#         # 读取医学图像
#         img_itk = sitk.ReadImage(file)
#         # 获取原始图像尺寸和间距
#         original_size = img_itk.GetSize()
#         original_spacing = img_itk.GetSpacing()
#         print("original_size:", original_size)
#         print("original_spacing:", original_spacing)
#
#         # 计算新的图像尺寸和间距
#         new_size = [512, 416, original_size[2]]
#         new_spacing = [original_size[0] * (original_spacing[0] / new_size[0]),
#                        original_size[1] * (original_spacing[1] / new_size[1]),
#                        original_size[2] * (original_spacing[2] / new_size[2])]
#         print("new_size:", new_size)
#         print("new_spacing:", new_spacing)
#
#         # 初始化重采样滤波器
#         resampleSliceFilter = sitk.ResampleImageFilter()
#         # 如果不是标签图像
#         if label == False:
#             resampleSliceFilter.SetOutputSpacing(new_spacing)
#             resampleSliceFilter.SetSize(new_size)
#             resampleSliceFilter.SetOutputDirection(img_itk.GetDirection())
#             resampleSliceFilter.SetOutputOrigin(img_itk.GetOrigin())
#             resampleSliceFilter.SetTransform(sitk.Transform())
#             resampleSliceFilter.SetDefaultPixelValue(img_itk.GetPixelIDValue())
#             resampleSliceFilter.SetInterpolator(sitk.sitkBSpline)
#             # 执行重采样
#             resampleImage = resampleSliceFilter.Execute(img_itk)
#             resampleImageArray = sitk.GetArrayFromImage(resampleImage)
#             resampleImageArray[resampleImageArray < 0] = 0  # 将图像中小于0的元素置为0
#         else:  # 对于标签，应使用线性插值以确保原始和重采样后的标签相同
#             resampleImage = resampleSliceFilter.Execute(img_itk)
#             resampleImageArray = sitk.GetArrayFromImage(resampleImage)
#
#         pad_height = (512 - new_size[1]) // 2  # 计算填充量
#         print(new_size[1], pad_height)
#         resampleImageArray_padded = np.pad(resampleImageArray, (
#             (0, 0), (pad_height, pad_height), (0, 0)
#         ), mode='constant', constant_values=0)  # 使用np.pad对图像进行上下填充
#
#         # 从NumPy数组创建SimpleITK图像
#         print(resampleImageArray_padded.shape)
#         img_itk_new = sitk.GetImageFromArray(resampleImageArray_padded.astype(np.float32))
#         img_itk_new.SetSpacing(new_spacing)
#         img_itk_new.SetDirection(img_itk.GetDirection())
#         img_itk_new.SetOrigin(img_itk.GetOrigin())
#
#         # 生成新的文件名并保存图像
#         # file_name = os.path.basename(file)
#         # file_name = file_name.replace('copy', 'resample')
#         # save_path = os.path.join(save_folder, file_name)
#         print("file_name:", output_path)
#
#         # 保存图像
#         sitk.WriteImage(img_itk_new, output_path)
#
#
# if test-main:
# input_path = "tumor-CONG PEI MIN_M00236746-13-t1_vibe_fs_tra_p4_bh_320_A.nii"
# output_path = "output-tumor-CONG PEI MIN_M00236746-13-t1_vibe_fs_tra_p4_bh_320_A.nii"
